import torch

class Imagenet:
    def __init__(self):
        from torchvision import models
        self.squeezenet = models.squeezenet1_1(pretrained=True)
        from torchvision import transforms
        self.transform = transforms.Compose([                #[1]
            transforms.Resize(256),                     #[2]
            transforms.CenterCrop(224),                 #[3]
            transforms.ToTensor(),                      #[4]
            transforms.Normalize(                       #[5]
            mean=[0.485, 0.456, 0.406],                 #[6]
            std=[0.229, 0.224, 0.225]                   #[7]
        )])
        self.squeezenet.eval()
        with open('resource/imagenet_classes.txt') as f:
            self.classes = [line.strip() for line in f.readlines()]

    def classify(self, img):
        # img = Image.open("jelly.jpeg")
        img_t = self.transform(img)
        batch_t = torch.unsqueeze(img_t, 0)
        out = self.squeezenet(batch_t)
        
        _, index = torch.max(out, 1)
        percentage = torch.nn.functional.softmax(out, dim=1)[0] * 100
        # 置信度 
        return self.classes[index[0]], percentage[index[0]].item()
